msg_tool\scripts\artemis\archive/
pfs.rs

1//! Artemis Engine PFS Archive (pf6 and pf8)
2use super::detect_script_type;
3use crate::ext::io::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::struct_pack::*;
7use anyhow::Result;
8use digest_io::IoWrapper;
9use msg_tool_macro::*;
10use sha1::Digest;
11use std::collections::HashMap;
12use std::io::{Read, Seek, SeekFrom, Write};
13use std::sync::{Arc, Mutex};
14
15#[derive(Debug)]
16/// The builder for Artemis PFS archive scripts.
17pub struct ArtemisArcBuilder {}
18
19impl ArtemisArcBuilder {
20    /// Creates a new instance of `ArtemisArcBuilder`.
21    pub fn new() -> Self {
22        ArtemisArcBuilder {}
23    }
24}
25
26impl ScriptBuilder for ArtemisArcBuilder {
27    fn default_encoding(&self) -> Encoding {
28        Encoding::Utf8
29    }
30
31    fn default_archive_encoding(&self) -> Option<Encoding> {
32        Some(Encoding::Utf8)
33    }
34
35    fn build_script(
36        &self,
37        buf: Vec<u8>,
38        filename: &str,
39        _encoding: Encoding,
40        archive_encoding: Encoding,
41        config: &ExtraConfig,
42        _archive: Option<&Box<dyn Script>>,
43    ) -> Result<Box<dyn Script + Send + Sync>> {
44        Ok(Box::new(ArtemisArc::new(
45            MemReader::new(buf),
46            archive_encoding,
47            config,
48            filename,
49        )?))
50    }
51
52    fn build_script_from_file(
53        &self,
54        filename: &str,
55        _encoding: Encoding,
56        archive_encoding: Encoding,
57        config: &ExtraConfig,
58        _archive: Option<&Box<dyn Script>>,
59    ) -> Result<Box<dyn Script + Send + Sync>> {
60        let f = std::fs::File::open(filename)?;
61        let f = std::io::BufReader::new(f);
62        Ok(Box::new(ArtemisArc::new(
63            f,
64            archive_encoding,
65            config,
66            filename,
67        )?))
68    }
69
70    fn build_script_from_reader<'a>(
71        &self,
72        reader: Box<dyn ReadSeek + Send + Sync + 'a>,
73        filename: &str,
74        _encoding: Encoding,
75        archive_encoding: Encoding,
76        config: &ExtraConfig,
77        _archive: Option<&Box<dyn Script>>,
78    ) -> Result<Box<dyn Script + Send + Sync + 'a>> {
79        Ok(Box::new(ArtemisArc::new(
80            reader,
81            archive_encoding,
82            config,
83            filename,
84        )?))
85    }
86
87    fn extensions(&self) -> &'static [&'static str] {
88        gen_artemis_arc_ext!()
89    }
90
91    fn is_archive(&self) -> bool {
92        true
93    }
94
95    fn script_type(&self) -> &'static ScriptType {
96        &ScriptType::ArtemisArc
97    }
98
99    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
100        if buf_len >= 3 && (buf.starts_with(b"pf6") || buf.starts_with(b"pf8")) {
101            return Some(10);
102        }
103        None
104    }
105
106    fn create_archive(
107        &self,
108        filename: &str,
109        files: &[&str],
110        encoding: Encoding,
111        config: &ExtraConfig,
112    ) -> Result<Box<dyn Archive>> {
113        let f = std::fs::File::options()
114            .write(true)
115            .read(true)
116            .create(true)
117            .truncate(true)
118            .open(filename)?;
119        Ok(Box::new(ArtemisArcWriter::new(f, files, encoding, config)?))
120    }
121}
122
123#[derive(Debug, Clone, StructPack, StructUnpack)]
124struct PfsEntryHeader {
125    #[pstring(u32)]
126    name: String,
127    _unk: u32,
128    offset: u32,
129    size: u32,
130}
131
132#[derive(Debug)]
133/// The Artemis PFS archive script.
134pub struct ArtemisArc<'a, T: Read + Seek + std::fmt::Debug + Send + Sync + 'a> {
135    reader: Arc<Mutex<T>>,
136    entries: Vec<PfsEntryHeader>,
137    xor_key: Option<[u8; 20]>,
138    output_ext: Option<String>,
139    _mark: std::marker::PhantomData<&'a ()>,
140}
141
142impl<'b, T: Read + Seek + std::fmt::Debug + Send + Sync + 'b> ArtemisArc<'b, T> {
143    /// Creates a new Artemis PFS archive script.
144    ///
145    /// * `reader` - The reader for the archive.
146    /// * `archive_encoding` - The encoding used for the archive.
147    /// * `config` - Extra configuration options.
148    /// * `filename` - The name of the archive file.
149    pub fn new(
150        mut reader: T,
151        archive_encoding: Encoding,
152        _config: &ExtraConfig,
153        filename: &str,
154    ) -> Result<Self> {
155        let mut magic = [0; 2];
156        reader.read_exact(&mut magic)?;
157        if &magic != b"pf" {
158            return Err(anyhow::anyhow!(
159                "Invalid Artemis archive magic: {:?}",
160                magic
161            ));
162        }
163        let version = reader.read_u8()?;
164        if version != b'6' && version != b'8' {
165            return Err(anyhow::anyhow!(
166                "Unsupported Artemis archive version: {}",
167                version
168            ));
169        }
170        let index_size = reader.read_u32()?;
171        let file_count = reader.read_u32()?;
172        let mut entries = Vec::with_capacity(file_count as usize);
173        for _ in 0..file_count {
174            let header = reader.read_struct(false, archive_encoding, &None)?;
175            entries.push(header);
176        }
177        let xor_key = if version == b'8' {
178            reader.seek(SeekFrom::Start(7))?;
179            let mut sha = IoWrapper(sha1::Sha1::default());
180            let ra = &mut reader;
181            let mut r = ra.take(index_size as u64);
182            std::io::copy(&mut r, &mut sha)?;
183            sha.flush()?;
184            let result = sha.0.finalize();
185            let mut xor_key = [0u8; 20];
186            xor_key.copy_from_slice(&result);
187            Some(xor_key)
188        } else {
189            None
190        };
191        let output_ext = std::path::Path::new(filename)
192            .extension()
193            .filter(|s| *s != "pfs")
194            .map(|s| s.to_string_lossy().to_string());
195        Ok(ArtemisArc {
196            reader: Arc::new(Mutex::new(reader)),
197            entries,
198            xor_key,
199            output_ext,
200            _mark: std::marker::PhantomData,
201        })
202    }
203}
204
205impl<'b, T: Read + Seek + std::fmt::Debug + Send + Sync + 'b> Script for ArtemisArc<'b, T> {
206    fn default_output_script_type(&self) -> OutputScriptType {
207        OutputScriptType::Json
208    }
209
210    fn default_format_type(&self) -> FormatOptions {
211        FormatOptions::None
212    }
213
214    fn is_archive(&self) -> bool {
215        true
216    }
217
218    fn iter_archive_filename<'a>(
219        &'a self,
220    ) -> Result<Box<dyn Iterator<Item = Result<String>> + 'a>> {
221        Ok(Box::new(
222            self.entries.iter().map(|header| Ok(header.name.clone())),
223        ))
224    }
225
226    fn iter_archive_offset<'a>(&'a self) -> Result<Box<dyn Iterator<Item = Result<u64>> + 'a>> {
227        Ok(Box::new(
228            self.entries.iter().map(|header| Ok(header.offset as u64)),
229        ))
230    }
231
232    fn open_file<'a>(&'a self, index: usize) -> Result<Box<dyn ArchiveContent + Send + Sync + 'a>> {
233        if index >= self.entries.len() {
234            return Err(anyhow::anyhow!(
235                "Index out of bounds: {} (max: {})",
236                index,
237                self.entries.len()
238            ));
239        }
240        let header = &self.entries[index];
241        let mut entry = Entry {
242            header: header.clone(),
243            reader: self.reader.clone(),
244            pos: 0,
245            script_type: None,
246            xor_key: self.xor_key.clone(),
247        };
248        let mut header = [0; 0x20];
249        let readed = entry.read(&mut header)?;
250        entry.pos = 0;
251        entry.script_type = detect_script_type(&header, readed, &entry.header.name);
252        Ok(Box::new(entry))
253    }
254
255    fn archive_output_ext<'a>(&'a self) -> Option<&'a str> {
256        self.output_ext.as_ref().map(|s| s.as_str())
257    }
258}
259
260#[derive(Debug)]
261struct Entry<T: Read + Seek + std::fmt::Debug> {
262    header: PfsEntryHeader,
263    reader: Arc<Mutex<T>>,
264    pos: u64,
265    script_type: Option<ScriptType>,
266    xor_key: Option<[u8; 20]>,
267}
268
269impl<T: Read + Seek + std::fmt::Debug + Send + Sync> ArchiveContent for Entry<T> {
270    fn name(&self) -> &str {
271        &self.header.name
272    }
273
274    fn size(&self) -> Option<u64> {
275        Some(self.header.size as u64)
276    }
277
278    fn script_type(&self) -> Option<&ScriptType> {
279        self.script_type.as_ref()
280    }
281
282    fn to_data<'a>(&'a mut self) -> Result<Box<dyn ReadSeek + Send + Sync + 'a>> {
283        Ok(Box::new(self))
284    }
285}
286
287impl<T: Read + Seek + std::fmt::Debug> Read for Entry<T> {
288    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
289        let mut reader = self.reader.lock().map_err(|e| {
290            std::io::Error::new(
291                std::io::ErrorKind::Other,
292                format!("Failed to lock mutex: {}", e),
293            )
294        })?;
295        reader.seek(SeekFrom::Start(self.header.offset as u64 + self.pos))?;
296        let bytes_read = buf.len().min(self.header.size as usize - self.pos as usize);
297        if bytes_read == 0 {
298            return Ok(0);
299        }
300        let bytes_read = reader.read(&mut buf[..bytes_read])?;
301        if let Some(xor_key) = &self.xor_key {
302            for i in 0..bytes_read {
303                let l = (self.pos + i as u64) % 20;
304                buf[i] ^= xor_key[l as usize];
305            }
306        }
307        self.pos += bytes_read as u64;
308        Ok(bytes_read)
309    }
310}
311
312impl<T: Read + Seek + std::fmt::Debug> Seek for Entry<T> {
313    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
314        let new_pos = match pos {
315            SeekFrom::Start(offset) => offset,
316            SeekFrom::End(offset) => {
317                if offset < 0 {
318                    if (-offset) as u64 > self.header.size as u64 {
319                        return Err(std::io::Error::new(
320                            std::io::ErrorKind::InvalidInput,
321                            "Seek from end exceeds file length",
322                        ));
323                    }
324                    self.header.size as u64 - (-offset) as u64
325                } else {
326                    self.header.size as u64 + offset as u64
327                }
328            }
329            SeekFrom::Current(offset) => {
330                if offset < 0 {
331                    if (-offset) as u64 > self.pos {
332                        return Err(std::io::Error::new(
333                            std::io::ErrorKind::InvalidInput,
334                            "Seek from current exceeds current position",
335                        ));
336                    }
337                    self.pos.saturating_sub((-offset) as u64)
338                } else {
339                    self.pos + offset as u64
340                }
341            }
342        };
343        self.pos = new_pos;
344        Ok(self.pos)
345    }
346
347    fn stream_position(&mut self) -> std::io::Result<u64> {
348        Ok(self.pos)
349    }
350}
351
352/// The Artemis PFS archive writer.
353pub struct ArtemisArcWriter<T: Write + Seek + Read> {
354    writer: T,
355    headers: HashMap<String, PfsEntryHeader>,
356    encoding: Encoding,
357    disable_xor: bool,
358    index_size: u32,
359}
360
361impl<T: Write + Seek + Read> ArtemisArcWriter<T> {
362    /// Creates a new Artemis PFS archive writer.
363    ///
364    /// * `writer` - The writer for the archive.
365    /// * `files` - The list of files to include in the archive.
366    /// * `encoding` - The encoding used for the archive.
367    /// * `config` - Extra configuration options.
368    pub fn new(
369        mut writer: T,
370        files: &[&str],
371        encoding: Encoding,
372        config: &ExtraConfig,
373    ) -> Result<Self> {
374        writer.write_all(if config.artemis_arc_disable_xor {
375            b"pf6"
376        } else {
377            b"pf8"
378        })?;
379        writer.write_u32(0)?; // Placeholder for index size
380        writer.write_u32(files.len() as u32)?;
381        let mut headers = HashMap::new();
382        for file in files {
383            let header = PfsEntryHeader {
384                name: file.to_string(),
385                _unk: 0,
386                offset: 0,
387                size: 0,
388            };
389            header.pack(&mut writer, false, encoding, &None)?;
390            headers.insert(file.to_string(), header);
391        }
392        let size = writer.stream_position()?;
393        let index_size = size as u32 - 7;
394        writer.write_u32_at(3, index_size)?;
395        Ok(ArtemisArcWriter {
396            writer,
397            headers,
398            encoding,
399            disable_xor: config.artemis_arc_disable_xor,
400            index_size,
401        })
402    }
403}
404
405impl<T: Write + Seek + Read> Archive for ArtemisArcWriter<T> {
406    fn new_file<'a>(
407        &'a mut self,
408        name: &str,
409        _size: Option<u64>,
410    ) -> Result<Box<dyn WriteSeek + 'a>> {
411        let entry = self
412            .headers
413            .get_mut(name)
414            .ok_or_else(|| anyhow::anyhow!("File '{}' not found in archive", name))?;
415        if entry.offset != 0 || entry.size != 0 {
416            return Err(anyhow::anyhow!("File '{}' already exists in archive", name));
417        }
418        self.writer.seek(SeekFrom::End(0))?;
419        entry.offset = self.writer.stream_position()? as u32;
420        let file = ArtemisArcFile {
421            header: entry,
422            writer: &mut self.writer,
423            pos: 0,
424        };
425        Ok(Box::new(file))
426    }
427
428    fn write_header(&mut self) -> Result<()> {
429        self.writer.seek(SeekFrom::Start(11))?;
430        let mut files = self.headers.values().collect::<Vec<_>>();
431        files.sort_by_key(|d| d.offset);
432        for file in files.iter() {
433            file.pack(&mut self.writer, false, self.encoding, &None)?;
434        }
435        if !self.disable_xor {
436            self.writer.seek(SeekFrom::Start(7))?;
437            let mut sha = IoWrapper(sha1::Sha1::default());
438            let w = &mut self.writer;
439            let mut header = w.take(self.index_size as u64);
440            std::io::copy(&mut header, &mut sha)?;
441            sha.flush()?;
442            let result = sha.0.finalize();
443            let mut xor_key = [0u8; 20];
444            xor_key.copy_from_slice(&result);
445            let mut buf = [0u8; 1024];
446            for file in files.iter() {
447                self.writer.seek(SeekFrom::Start(file.offset as u64))?;
448                let mut pos = 0u32;
449                while pos < file.size {
450                    let bytes_to_read = (file.size - pos).min(1024) as usize;
451                    let bytes_read = self.writer.read(&mut buf[..bytes_to_read])?;
452                    if bytes_read == 0 {
453                        return Err(anyhow::anyhow!(
454                            "Unexpected end of file while reading '{}'",
455                            file.name
456                        ));
457                    }
458                    for i in 0..bytes_read {
459                        let l = (pos as u64 + i as u64) % 20;
460                        buf[i] ^= xor_key[l as usize];
461                    }
462                    self.writer.seek_relative(-(bytes_read as i64))?;
463                    self.writer.write_all(&buf[..bytes_read])?;
464                    pos += bytes_read as u32;
465                }
466            }
467        }
468        Ok(())
469    }
470}
471
472/// The Artemis PFS archive file writer.
473pub struct ArtemisArcFile<'a, T: Write + Seek> {
474    header: &'a mut PfsEntryHeader,
475    writer: &'a mut T,
476    pos: u64,
477}
478
479impl<'a, T: Write + Seek> Write for ArtemisArcFile<'a, T> {
480    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
481        self.writer
482            .seek(SeekFrom::Start(self.header.offset as u64 + self.pos))?;
483        let bytes_written = self.writer.write(buf)?;
484        self.pos += bytes_written as u64;
485        self.header.size = self.header.size.max(self.pos as u32);
486        Ok(bytes_written)
487    }
488
489    fn flush(&mut self) -> std::io::Result<()> {
490        self.writer.flush()
491    }
492}
493
494impl<'a, T: Write + Seek> Seek for ArtemisArcFile<'a, T> {
495    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
496        let new_pos = match pos {
497            SeekFrom::Start(offset) => offset,
498            SeekFrom::End(offset) => {
499                if offset < 0 {
500                    if (-offset) as u64 > self.header.size as u64 {
501                        return Err(std::io::Error::new(
502                            std::io::ErrorKind::InvalidInput,
503                            "Seek from end exceeds file length",
504                        ));
505                    }
506                    self.header.size as u64 - (-offset) as u64
507                } else {
508                    self.header.size as u64 + offset as u64
509                }
510            }
511            SeekFrom::Current(offset) => {
512                if offset < 0 {
513                    if (-offset) as u64 > self.pos {
514                        return Err(std::io::Error::new(
515                            std::io::ErrorKind::InvalidInput,
516                            "Seek from current exceeds current position",
517                        ));
518                    }
519                    self.pos.saturating_sub((-offset) as u64)
520                } else {
521                    self.pos + offset as u64
522                }
523            }
524        };
525        self.pos = new_pos;
526        Ok(self.pos)
527    }
528
529    fn stream_position(&mut self) -> std::io::Result<u64> {
530        Ok(self.pos)
531    }
532}